{-# LANGUAGE TypeFamilies, MultiParamTypeClasses #-}

module LambdaLift where

import Common
import qualified AnonymousFunctions as A
import AnonymousFunctions
  (AnonTerm, AnonTerm(Lam), Lam_(..), Let_(..), Decl(..))
import qualified TopLevelFunctions as T
import TopLevelFunctions (TopTerm, TopProg(..))

import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

import Yoko
import CnvExample2A

import LLBasics
import FreeVars (bump, anonFVs)
import DeepSeq (DeepSeq(..))





-- 28 lines (~6 of which are specific to Yoko)
lambdaLift :: [Type] -> AnonTerm -> TopProg
lambdaLift e x = Prog ds tm where
  (tm, ds) = runMnd_LL (ll x) (e, (length e, IM.empty), 0)

ll x = convert Cnv_LL x
data Cnv_LL = Cnv_LL
type instance Idiom Cnv_LL = Mnd_LL



type instance Cnv Cnv_LL AnonTerm = TopTerm
instance Convert Cnv_LL AnonTerm where
  convert cnv tm = case partition $ disband tm of
    Left x -> ($ x) $    llLam .|| llVar .|. llLet
    Right x -> convert cnv x

llLam (Lam_ tyTop tmTop) = do
  let ((tys, ty), tm) = peel ([], tyTop) tmTop where
        peel (acc, ty') (A.Lam ty tm) = peel (ty' : acc, ty) tm
        peel acc tm = (acc, tm)

  let nLocals = 1 + length tys -- NB "1 +" is for ty

  let captured = bump nLocals $ anonFVs tm
      captured' = IS.toAscList captured
      captured'' = reverse captured'

  (rho, rn) <- ask

  (let updE _ = (,) (ty : tys ++ rho) . (,) nLocals .
                IM.fromDistinctAscList $ zip captured' [0..]
   in local updE (ll tm)) >>= \tm ->
      emit (map (rho !!) captured'', reverse tys, ty, tm)

  sh <- numEmissions

  return . T.DVar (sh - 1) . map (lookupRN rn) $ captured''
llVar (A.Var_ i) = ask >>= \(_, rn) -> return . T.Var $ lookupRN rn i
llLet (Let_ ds tm) = ll $ foldr (\(Decl ty tm) x -> A.App (Lam ty tm) x) tm ds





infixl 1 @@
(@@) = A.App

s_comb a b c = Lam (TyFun a (TyFun b c)) . Lam (TyFun a b) . Lam a $
               A.Var 2 @@ A.Var 0 @@ (A.Var 1 @@ A.Var 0)

ex0 = Lam TyInt (A.Var 0)
ex0' = lambdaLift [] ex0

ex1 = s_comb TyInt TyInt TyInt @@ (Lam TyInt $ Lam TyInt (A.Var 0))
                               @@ Lam TyUnit (A.Var 2 @@ A.Var 1)
ex1' = lambdaLift [TyInt, TyFun TyInt TyInt] ex1

ex2 = Lam (TyFun (TyFun TyInt TyInt) TyUnit) $
      Lam (TyFun TyInt TyInt)
            (A.Var 1 @@ Lam TyInt (A.Var 1 @@ A.Var 0))
ex2' = lambdaLift [] ex2

ex3 = Lam TyUnit . Lam TyUnit . Lam TyUnit .
      (A.Var 1 @@) . Lam TyUnit .
       (A.Var 3 @@) . Lam TyInt $
        A.Var 1
ex3' = lambdaLift [] ex3

ex4 = Lam (TyFun TyInt TyInt) (A.Var 0) @@ Lam TyInt (A.Var 0)
ex4' = lambdaLift [] ex4

ex5 = Lam (TyFun (TyFun TyInt TyInt) (TyFun TyInt TyInt)) (A.Var 0) @@
      Lam (TyFun TyInt TyInt)
          (Lam TyUnit (A.Var 1) @@ A.Var 1)
ex5' = lambdaLift [TyUnit] ex5

-- TODO can I make the DVars use scoped instead of global indices?

{-
ex0' ==
Prog [([],[],TyInt,Var 0)
     ] (DVar 0 [])

ex1' ==
Prog [([],[TyFun TyInt (TyFun TyInt TyInt),TyFun TyInt TyInt],TyInt,
         App (App (Var 2) (Var 0)) (App (Var 1) (Var 0))),
      ([],[TyInt],TyInt,Var 0),
      ([TyFun TyInt TyInt,TyInt],[],TyUnit,App (Var ^1) (Var ^0))
     ] (App (App (DVar 0 []) (DVar 1 [])) (DVar 2 [1,0]))

ex2' ==
Prog [([TyFun TyInt TyInt],[],TyInt,App (Var ^0) (Var 0)),
      ([],[TyFun (TyFun TyInt TyInt) TyUnit],TyFun TyInt TyInt,
         App (Var 1) (DVar 0 [0]))
     ] (DVar 1 [])

ex3' ==
Prog [([TyUnit],[],TyInt,Var ^0),
      ([TyUnit],[],TyUnit,App (Var ^0) (DVar 0 [0])),
      ([],[TyUnit,TyUnit],TyUnit,App (Var 1) (DVar 1 [2]))
     ] (DVar 2 [])

ex4' ==
Prog [([],[],TyFun TyInt TyInt,Var 0),
      ([],[],TyInt,Var 0)
     ] (App (DVar 0 []) (DVar 1 []))

ex5' ==
Prog [([],[],TyFun (TyFun TyInt TyInt) (TyFun TyInt TyInt),Var 0),
      ([TyFun TyInt TyInt],[],TyUnit,Var ^0),
      ([TyUnit],[],TyFun TyInt TyInt,App (DVar 1 [0]) (Var ^0))
     ] (App (DVar 0 []) (DVar 2 [0]))

-}

instance DeepSeq Type    where rnf = (`seq` ())
instance DeepSeq T.Occ   where rnf (T.Occ x) = rnf x
instance DeepSeq TopProg where rnf (Prog decs tm) = rnf decs `seq` rnf tm
instance DeepSeq TopTerm where rnf = rnf . reps . disband

all_exs = rnf [ex0', ex1', ex2', ex3', ex4', ex5']
